package Core.plugin

import Core.{CPUSimple, Plugin}
import spinal.core._
import spinal.lib._

case class LLbitReadPorts()  extends Bundle with IMasterSlave{
    val LLbit_data = Bool()

    override def asMaster(): Unit = {
        in(LLbit_data)
    }
}

case class LLbitWritePorts() extends Bundle with IMasterSlave{
    val LLbit_wen = Bool()
    val LLbit_data = Bool()

    override def asMaster(): Unit = {
        out(LLbit_wen, LLbit_data)
    }
}

case class LLbitModule() extends Component{
    val read_ports = slave(LLbitReadPorts())
    val write_ports = slave(LLbitWritePorts())
    val LLbit = RegInit(False)

    when(write_ports.LLbit_wen) {
        LLbit := write_ports.LLbit_data
    }

    read_ports.LLbit_data := write_ports.LLbit_wen ? write_ports.LLbit_data | LLbit
}

class LsuPlugin(AW : Int = 32, DW : Int = 32) extends Plugin[CPUSimple] with DCacheAccessService {
    @dontName var dcache_access : DCacheAccess = null
    override def newDCacheAccess(): DCacheAccess = {
        assert(dcache_access == null)
        dcache_access = DCacheAccess(AW, DW)
        dcache_access
    }

    override def setup(pipeline: CPUSimple): Unit = {

    }

    override def build(pipeline: CPUSimple): Unit = {
        import pipeline._
        import pipeline.config._

        val LLbit_module = new LLbitModule()
        memaccess plug new Area {
            import memaccess._
            val cpu_addr = input(ALU_RESULT).asUInt
            val is_mem = input(IS_STORE) || input(IS_LOAD)
            val cpu_addr_sel = cpu_addr(1 downto 0)
            val lsu_addr  = UInt(XLEN bits)
            val lsu_rdata = Bits(XLEN bits)
            val lsu_wdata = Bits(XLEN bits)
            val lsu_wen   = Bool()
            val lsu_wstrb = Bits(XLEN/8 bits)
            val lsu_size  = UInt(2 bits)

            val dcache_rdata = dcache_access.rsp.payload.data |>> (cpu_addr_sel << 3)
            val dcache_ld_b   = B((XLEN-8-1 downto 0) -> dcache_rdata(7)) ## dcache_rdata(7 downto 0)
            val dcache_ld_bu  = B((XLEN-8-1 downto 0) -> False) ## dcache_rdata(7 downto 0)
            val dcache_ld_h   = B((XLEN-16-1 downto 0) -> dcache_rdata(15)) ## dcache_rdata(15 downto 0)
            val dcache_ld_hu  = B((XLEN-16-1 downto 0) -> False) ## dcache_rdata(15 downto 0)
            val dcache_ld_w   = dcache_rdata(31 downto 0)
            val dcache_data_load = Bits(XLEN bits)

            val mem_wdata = input(MEM_WDATA)
            val dcache_st_b = B((XLEN-8-1 downto 0) -> mem_wdata(7)) ## mem_wdata(7 downto 0)
            val dcache_st_h = B((XLEN-16-1 downto 0) -> mem_wdata(15)) ## mem_wdata(15 downto 0)
            val dcache_st_w = mem_wdata(31 downto 0)
            val dcache_wdata = Bits(XLEN bits)
            val dcache_wstrb = Bits(XLEN/8 bits)
            val dcache_wen   = input(IS_STORE)

            val is_ale = Bool()

            insert(LLBIT) := LLbit_module.read_ports.LLbit_data
            val LLbit = input(LLBIT)
            val LLbit_data = False
            val LLbit_we = False
            val is_sc = input(MEM_CTRL) === MemCtrlEnum.LL.asBits
            val is_ll = input(MEM_CTRL) === MemCtrlEnum.SC.asBits

            when(is_sc) {
                when(LLbit === False) {
                    dcache_wen := False
                } .otherwise {
                    LLbit_data := False
                    LLbit_we := True
                }
            }

            when(is_ll) {
                LLbit_data := True
                LLbit_we := False
            }

            is_ale := input(MEM_CTRL).mux(
                MemCtrlEnum.LD_H.asBits -> (cpu_addr.asBits(0) =/= False),
                MemCtrlEnum.LD_HU.asBits -> (cpu_addr.asBits(0) =/= False),
                MemCtrlEnum.LD_W.asBits -> (cpu_addr.asBits(1 downto 0) =/= B(0, 2 bits)),
                MemCtrlEnum.LL.asBits -> (cpu_addr.asBits(1 downto 0) =/= B(0, 2 bits)),
                MemCtrlEnum.ST_H.asBits -> (cpu_addr.asBits(0) =/= False),
                MemCtrlEnum.ST_W.asBits -> (cpu_addr.asBits(1 downto 0) =/= B(0, 2 bits)),
                MemCtrlEnum.SC.asBits -> (cpu_addr.asBits(1 downto 0) =/= B(0, 2 bits)),
                default -> False
            )

            insert(IS_ALE) := is_ale

            dcache_data_load := input(MEM_CTRL).mux(
                MemCtrlEnum.LD_B.asBits -> dcache_ld_b,
                MemCtrlEnum.LD_BU.asBits -> dcache_ld_bu,
                MemCtrlEnum.LD_H.asBits -> dcache_ld_h,
                MemCtrlEnum.LD_HU.asBits -> dcache_ld_hu,
                MemCtrlEnum.LD_W.asBits -> dcache_ld_w,
                MemCtrlEnum.LL.asBits -> dcache_ld_w,
                MemCtrlEnum.ST_B.asBits -> dcache_st_b,
                MemCtrlEnum.ST_H.asBits -> dcache_st_h,
                MemCtrlEnum.ST_W.asBits -> dcache_st_w,
                MemCtrlEnum.SC.asBits -> dcache_st_w,
                default -> B(0, XLEN bits)
            )

            lsu_size := input(MEM_CTRL).mux(
                MemCtrlEnum.LD_B.asBits -> U(0),
                MemCtrlEnum.LD_BU.asBits -> U(0),
                MemCtrlEnum.LD_H.asBits -> U(1),
                MemCtrlEnum.LD_HU.asBits -> U(1),
                MemCtrlEnum.LD_W.asBits -> U(2),
                MemCtrlEnum.LL.asBits -> U(2),
                default -> U(0)
            )

            dcache_wdata := input(MEM_CTRL).mux(
                MemCtrlEnum.ST_B.asBits -> dcache_st_b,
                MemCtrlEnum.ST_H.asBits -> dcache_st_h,
                MemCtrlEnum.ST_W.asBits -> dcache_st_w,
                MemCtrlEnum.SC.asBits -> dcache_st_w,
                default -> B(0, XLEN bits)
            )

            dcache_wstrb := input(MEM_CTRL).mux(
                MemCtrlEnum.ST_B.asBits -> B(4 bits, 0 -> True, default -> False),
                MemCtrlEnum.ST_H.asBits -> B(4 bits, (1 downto 0) -> True, default -> False),
                MemCtrlEnum.ST_W.asBits -> B(4 bits, (3 downto 0) -> True),
                MemCtrlEnum.SC.asBits   -> B(4 bits, (3 downto 0) -> True),
                default -> B(0, 4 bits)
            )

            insert(LLBIT_WE) := LLbit_we
            insert(LLBIT_DATA) := LLbit_data

            lsu_rdata := dcache_data_load  //TODO:need add unalign access
            lsu_wdata := dcache_wdata |<< (cpu_addr_sel << 3)//TODO:need add unalign access
            lsu_addr  := cpu_addr //TODO:need add unalign access
            lsu_wen   := dcache_wen
            lsu_wstrb := dcache_wstrb |<< cpu_addr_sel

            insert(LSU_RDATA):= lsu_rdata
            insert(LSU_WDATA):= lsu_wdata
            insert(LSU_HOLD) := dcache_access.stall
            insert(RD) := input(IS_LOAD) ? input(LSU_RDATA) | (input(IS_CSR) ? input(CSR_RDATA) | input(ALU_RESULT))

            dcache_access.cmd.valid        := is_mem && arbitration.isValid
            dcache_access.cmd.payload.addr := lsu_addr.resized
            dcache_access.cmd.payload.wen  := lsu_wen
            dcache_access.cmd.payload.wdata:= lsu_wdata
            dcache_access.cmd.payload.wstrb:= lsu_wstrb
            dcache_access.cmd.payload.size := lsu_size
        }

        writeback plug new Area{
            import writeback._
            //insert(RD) := input(IS_LOAD) ? input(LSU_RDATA) | input(ALU_RESULT)
            LLbit_module.write_ports.LLbit_data := input(LLBIT_DATA)
            LLbit_module.write_ports.LLbit_wen := input(LLBIT_WE)
            // for test:
            output(PC) := input(PC)
            output(INSTRUCTION) := input(INSTRUCTION)
        }
    }
}
